# 导入必要的库和模块
import os
import json
from torchvision.ops import box_iou
import torch

# 定义原始数据的JSON文件路径
origin_json_path = (
    "/home/jxy/datasets/gen-vlkt/hico_20160224_det/annotations/trainval_hico.json"
)

# 定义生成数据的文件夹路径
generate_hico_dir = "/home/jxy/program/hoi2/600HOI-600HOI_no_object-add_seg_only_instance-max_adapt-dynamic_threshold_add_weight_23-black_add_correlation_hoi/results/weakly0915save_hoi/add_hoi_labels_1p1o_0p0o_0920"

# 打开原始JSON文件并读取数据
with open(origin_json_path, "r", encoding="utf8") as f:
    json_data = json.load(f)

    # 初始化一些统计计数器
    count_temp = 0
    origin_total_hoi_count = 0
    origin_total_boxes_count = 0
    generate_total_hoi_count = 0
    generate_total_boxes_count = 0
    correct_match_hoi_count = 0
    correct_match_boxes_count = 0
    all_include_hoi_count = 0
    ge_half_include_hoi_count = 0
    l_half_include_hoi_count = 0
    zero_include_hoi_count = 0
    total_image_count = 0
    count_ = 1

    # 开始遍历原始数据中的每个图像
    for hoi in json_data:
        # 控制遍历的图像数量，防止处理太多图像
        if count_ > 1000000:
            break
        else:
            count_ += 1

        # 输出当前处理的图像文件名
        print(hoi["file_name"])

        # 统计图像总数
        total_image_count += 1

        # 获取图像文件名并构建对应的txt文件名
        file_name = hoi["file_name"]
        txt_name = file_name.split(".")[0] + ".txt"

        # 打开生成数据的txt文件并读取生成的HOI数据
        with open(os.path.join(generate_hico_dir, txt_name), "r") as f:
            generate_hoi_dict = eval(f.read())

        # 获取原始数据中的物体框和HOI数据
        origin_boxes = hoi["annotations"]
        origin_hoi = hoi["hoi_annotation"]

        # 获取生成的数据中的物体框和HOI数据
        generate_boxes = generate_hoi_dict["annotations"]
        generate_hoi = generate_hoi_dict["hoi_annotation"]

        # 统计原始数据和生成数据中的HOI和物体框数量
        origin_total_hoi_count += len(origin_hoi)
        generate_total_hoi_count += len(generate_hoi)
        origin_total_boxes_count += len(origin_boxes)
        generate_total_boxes_count += len(generate_boxes)

        # 初始化一个字典用于记录原始数据中的物体框与生成数据中的物体框的匹配关系
        origin2generate_index_dict = {}

        # 开始匹配原始数据中的物体框与生成数据中的物体框
        for i in range(len(origin_boxes)):
            box_i = origin_boxes[i]["bbox"]
            max_iou = -1
            max_j = -1
            for j in range(len(generate_boxes)):
                box_j = generate_boxes[j]["bbox"]
                iou_ = box_iou(
                    torch.tensor(box_i).unsqueeze(0), torch.tensor(box_j).unsqueeze(0)
                )
                # 计算IoU，寻找最匹配的生成物体框
                if max_iou < iou_:
                    max_iou = iou_
                    max_j = j
            # 如果找到匹配的物体框且类别相同，则记录匹配关系
            if (
                max_iou > 0.5
                and origin_boxes[i]["category_id"]
                == generate_boxes[max_j]["category_id"]
            ):
                origin2generate_index_dict[i] = max_j
                correct_match_boxes_count += 1
            else:
                origin2generate_index_dict[i] = -1

        # 初始化一个计数器，用于统计正确匹配的HOI数量
        current_find_count = 0

        # 开始匹配原始HOI与生成HOI
        for i in range(len(origin_hoi)):
            hoi_i = origin_hoi[i]
            hoi_i_subject_id = hoi_i["subject_id"]
            hoi_i_object_id = hoi_i["object_id"]
            hoi_i_category_id = hoi_i["category_id"]
            hoi_i_hoi_category_id = hoi_i["hoi_category_id"]
            if (
                origin2generate_index_dict[hoi_i_subject_id] == -1
                or origin2generate_index_dict[hoi_i_object_id] == -1
            ):
                continue
            flag = False
            for j in range(len(generate_hoi)):
                hoi_j = generate_hoi[j]
                hoi_j_subject_id = hoi_j["subject_id"]
                hoi_j_object_id = hoi_j["object_id"]
                hoi_j_category_id = hoi_j["category_id"]
                hoi_j_hoi_category_id = hoi_j["hoi_category_id"]
                if (
                    origin2generate_index_dict[hoi_i_subject_id] == hoi_j_subject_id
                    and origin2generate_index_dict[hoi_i_object_id] == hoi_j_object_id
                    and hoi_i_category_id == hoi_j_category_id
                    and hoi_i_hoi_category_id == hoi_j_hoi_category_id
                ):
                    # 如果匹配成功，增加正确匹配的HOI数量
                    correct_match_hoi_count += 1
                    current_find_count += 1
                    flag = True
                    break

        # 根据匹配结果统计不同类型的HOI数量
        if current_find_count == 0:
            zero_include_hoi_count += 1
        elif current_find_count == len(origin_hoi):
            all_include_hoi_count += 1
        elif current_find_count / len(origin_hoi) >= 0.5:
            ge_half_include_hoi_count += 1
        else:
            l_half_include_hoi_count += 1

# 打开一个文件，将统计结果写入文件中
with open("add_hoi_labels_1p1o_0p0o_0920.txt", "w") as file:
    # 写入各种统计信息
    file.write("=========================================================\n")
    file.write("共%d张图像\n" % total_image_count)
    file.write("=========================================================\n")
    file.write("ground_truth所有关系数: %d\n" % origin_total_hoi_count)
    file.write("generate_ground_truth所有关系数: %d\n" % generate_total_hoi_count)
    file.write(
        "生成正确关系的数量: %d, 占原比: %.2f, 占生成比: %.2f\n"
        % (
            correct_match_hoi_count,
            correct_match_hoi_count / origin_total_hoi_count,
            correct_match_hoi_count / generate_total_hoi_count,
        )
    )
    file.write(
        "其中关系全部包括的: %d, 占比: %.2f\n"
        % (all_include_hoi_count, all_include_hoi_count / total_image_count)
    )
    file.write(
        "其中关系包括大于等于一半的: %d, 占比: %.2f\n"
        % (ge_half_include_hoi_count, ge_half_include_hoi_count / total_image_count)
    )
    file.write(
        "其中关系包括小于一半的: %d, 占比: %.2f\n"
        % (l_half_include_hoi_count, l_half_include_hoi_count / total_image_count)
    )
    file.write(
        "一个都不对的: %d, 占比: %.2f\n"
        % (zero_include_hoi_count, zero_include_hoi_count / total_image_count)
    )
    file.write("=========================================================\n")
    file.write("ground_truth所有检测框数: %d\n" % origin_total_boxes_count)
    file.write("generate_ground_truth所有检测框数: %d\n" % generate_total_boxes_count)
    file.write(
        "正确检测框个数: %d, 占原比: %.2f, 占生成比: %.2f\n"
        % (
            correct_match_boxes_count,
            correct_match_boxes_count / origin_total_boxes_count,
            correct_match_boxes_count / generate_total_boxes_count,
        )
    )
    file.write("=========================================================\n")
